#include "modular_reduction.h"
#include <cstdint>
#include <vector>

/*
 * @brief 计算Barrett约简所需的参数。
 * @param q 模数，64位无符号整数。
 * @return 包含三个参数的向量：
 * - mhi：Barrett约简的高位参数
 * - mlo：Barrett约简的低位参数
 * - 2*q：模数的两倍
 */
std::vector<uint64_t> BRedParams(uint64_t q)
{
    mpz_class bigR = NewUint(1);
    bigR <<= 128;
    bigR /= NewUint(q);
    mpz_class bigR_shifted = bigR >> 64;
    uint64_t mhi = mpz_get_ui(bigR_shifted.get_mpz_t());
    uint64_t mlo = mpz_get_ui(bigR.get_mpz_t());
    return { mhi, mlo, 2 * q };
}
/*
 * @brief Barrett 加法模运算。
 *
 * @param a 输入数值。
 * @param q 模数。
 * @param u Barrett 约简参数。
 *
 * @return 加法模运算结果。
 */
uint64_t BRedAdd(uint64_t a, uint64_t q, const std::vector<uint64_t> &u)
{
    uint64_t mhi, r, tmp;
    Mul64(a, u[0], mhi, tmp);
    r = a - mhi * q;
    if (r >= q) {
        r -= q;
    }
    return r;
}
/*
 * @brief Barrett 乘法模运算。
 *
 * 该函数计算并返回 Barrett 乘法模运算结果。通过 Barrett 约简，优化了乘法与模运算的速度，特别适合大数运算中的加速。
 *
 * @param x 输入数值。
 * @param y 输入数值。
 * @param q 模数。
 * @param u Barrett 约简参数。
 *
 * @return Barrett 乘法模运算结果。
 */
uint64_t BRed(uint64_t x, uint64_t y, uint64_t q, const std::vector<uint64_t> &u)
{
    uint64_t mhi, mlo, lhi, hhi, hlo, s0, carry, r, tmp;
    Mul64(x, y, mhi, mlo);
    r = mhi * u[0];
    Mul64(mlo, u[0], hhi, hlo);
    r += hhi;
    Mul64(mlo, u[1], lhi, tmp);
    Add64(hlo, lhi, 0, s0, carry);
    r += carry;
    Mul64(mhi, u[1], hhi, hlo); // mhi * ulo
    r += hhi;
    Add64(hlo, s0, 0, tmp, carry);
    r += carry;
    r = mlo - r * q;
    if (r >= q) {
        r -= q;
    }
    return r;
}
/*
 * @brief 常规模运算。
 *
 * 该函数计算 `a mod q`。
 *
 * @param a 输入数值。
 * @param q 模数。
 *
 * @return 模运算结果。
 */
uint64_t CRed(uint64_t a, uint64_t q)
{
    if (a >= q) {
        return a - q;
    }
    return a;
}
/*
 * @brief 64位整数乘法并分割高低位。
 *
 * 该函数执行两个64位整数的乘法，并将结果分割成高低位返回。计算 `x * y` 的高位和低位部分。
 *
 * @param x 第一个乘数。
 * @param y 第二个乘数。
 * @param mhi 结果的高位部分。
 * @param mlo 结果的低位部分。
 */
void Mul64(uint64_t x, uint64_t y, uint64_t &mhi, uint64_t &mlo)
{
    const uint64_t mask32 = (1ULL << 32) - 1;
    uint64_t x0 = x & mask32;
    uint64_t x1 = x >> 32;
    uint64_t y0 = y & mask32;
    uint64_t y1 = y >> 32;

    uint64_t w0 = x0 * y0;
    uint64_t t = x1 * y0 + (w0 >> 32);
    uint64_t w1 = t & mask32;
    uint64_t w2 = t >> 32;

    w1 += x0 * y1;
    uint64_t hi = x1 * y1 + w2 + (w1 >> 32);
    uint64_t lo = x * y;

    mhi = hi, mlo = lo;
}
/*
 * @brief 64位加法运算并处理进位。
 *
 * 该函数执行64位整数加法，考虑到进位的情况，将进位存储在 `carryOut` 中。
 *
 * @param x 第一个加数。
 * @param y 第二个加数。
 * @param carry 进位。
 * @param sum 加法结果。
 * @param carryOut 输出进位。
 */
void Add64(uint64_t x, uint64_t y, uint64_t carry, uint64_t &sum, uint64_t &carryOut)
{
    sum = x + y + carry;
    carryOut = ((x & y) | ((x | y) & (~sum))) >> 63;
}
/*
 * @brief 对向量进行标量乘法并应用 Barrett 约简。
 *
 * 该函数对向量 `p1` 中的每个元素与标量 `scalarMont` 进行乘法计算，并将结果通过 Barrett 约简存储到 `p2` 中。
 *
 * @param p1 输入向量，包含待乘的元素。
 * @param p2 输出向量，存储乘法结果。
 * @param scalarMont 乘法的标量，使用 Mont 息格式。
 * @param qi 模数，用于 Barrett 约简。
 * @param len 向量长度。
 */
void MulScalarBarrettVec(std::vector<uint64_t> &p1, std::vector<uint64_t> &p2, uint64_t scalarMont, uint64_t qi,
    uint64_t len)
{
    std::vector<uint64_t> mBredParams = BRedParams(qi);
    for (uint64_t j = 0; j < len; j++) {
        MulMod(p2[j], p1[j], scalarMont, qi);
    }
}
/*
 * @brief 对两个向量执行 Barrett 乘法，并将结果存储到第三个向量。
 *
 * 该函数对输入向量 `p1` 和 `p2` 中的每个元素执行 Barrett 乘法，并将结果存储到 `p3` 向量中。
 *
 * @param p1 第一个输入向量，包含待乘的系数。
 * @param p2 第二个输入向量，包含待乘的系数。
 * @param p3 输出向量，存储乘法结果。
 * @param qi 模数，用于 Barrett 约简。
 * @param len 向量长度。
 */
void MulCoeffsBarrettConstantVec(std::vector<uint64_t> &p1, std::vector<uint64_t> &p2, std::vector<uint64_t> &p3,
    uint64_t qi, uint64_t len)
{
    std::vector<uint64_t> mBredParams = BRedParams(qi);
    #pragma omp parallel for
    for (uint64_t j = 0; j < len; j++) {
        MulModBarrett(p3[j], p1[j], p2[j], qi, mBredParams[2], mBredParams[0], mBredParams[1]);
    }
}
/*
 * @brief 对两个向量中的系数执行 Barrett 乘法并将结果加到第三个向量中。
 *
 * 该函数对输入的两个向量 `p1` 和 `p2` 的每个元素执行 Barrett 乘法，并将结果累加到 `p3` 向量的相应元素中。
 * 乘法结果不进行模操作，而是直接进行加法。
 *
 * @param p1 输入向量，包含待乘的系数。
 * @param p2 输入向量，包含待乘的系数。
 * @param p3 输出向量，存储乘法结果的累加值。
 * @param qi 模数，用于 Barrett 乘法的计算。
 * @param len 向量的长度，即 `p1`、`p2`、`p3` 的元素个数。
 */
void MulCoeffsBarrettConstantAndAddNoModVec(std::vector<uint64_t> &p1, std::vector<uint64_t> &p2,
    std::vector<uint64_t> &p3, uint64_t qi, uint64_t len)
{
    std::vector<uint64_t> mBredParams = BRedParams(qi);
    #pragma omp parallel for
    for (uint64_t j = 0; j < len; j++) {
        uint64_t tmp;
        MulModBarrett(tmp, p1[j], p2[j], qi, mBredParams[2], mBredParams[0], mBredParams[1]);
        p3[j] += tmp;
    }
}
